package com.gitee.wsl.ext.map

import kotlin.collections.set

typealias DoubleKeyMap<Key1, Key2, Value> = MutableMap<Key1, MutableMap<Key2, Value>>

fun<Key1, Key2, Value> DoubleKeyMap<Key1, Key2, Value>.put(k1: Key1, k2: Key2, v: Value): Value? {
    var data2 = get(k1)
    var prev: Value? = null

    if (data2 == null) {
        data2 = LinkedHashMap()
        this[k1] = data2
    } else {
        prev = data2[k2]
    }

    data2[k2] = v
    return prev
}

operator fun<Key1, Key2, Value> DoubleKeyMap<Key1, Key2, Value>.get(k1: Key1, k2: Key2): Value? {
    val data2 = this[k1] ?: return null
    return data2[k2]
}

operator fun<Key1, Key2, Value> DoubleKeyMap<Key1, Key2, Value>.get(k1: Key1): Map<Key2, Value> =
    this.getOrPut(k1){ LinkedHashMap() }


/**
 * Get all values associated with primary key.
 */
fun<Key1, Key2, Value> DoubleKeyMap<Key1, Key2, Value>.values(k1: Key1): Collection<Value>? {
    val data2 = this[k1] ?: return null
    return data2.values
}

/**
 * Get all secondary keys associated with a primary key.
 */
 fun<Key1, Key2, Value> DoubleKeyMap<Key1, Key2, Value>.keySet(k1: Key1): Set<Key2>? {
    val data2 = this[k1] ?: return null
    return data2.keys
}